package main

import (
	"embed"
	_ "embed"
	"flag"
	"fmt"
	"html/template"
	"io"
	"os"
	"strconv"
)

/*
anchors: predetermined set of boxes with particular height-width ratios
mask: list of ids of the bounding boxes(anchors) that the layer is responsible for predicting
num: total number of anchors
filter = (num_classes+5)*k where k=number of masks in one yolo layer

mask = 0,1,2
anchors = 12,16, 19,36, 40,28, 36,75, 76,55, 72,146, 142,110, 192,243, 459,401 //高宽比
classes=3
num=9
*/
type YoloV4Param struct {
	Batch       int
	Subdivision int
	Classes     int
	MaxBatches  int
	Steps       [2]int
	Width       int
	Height      int
	Filters     int
}

func NewDefaultYoloV4Param() YoloV4Param {
	param := YoloV4Param{
		Batch:       64,
		Subdivision: 16,
		Classes:     3,
		MaxBatches:  6000,               //auto
		Steps:       [2]int{4800, 5400}, //auto
		Width:       416,
		Height:      416,
		Filters:     255, //auto
	}
	return param
}

// 根据输入参数，更新相关参数值
func (param *YoloV4Param) Update() {
	param.MaxBatches = param.Classes * 2000
	param.Steps[0] = int(float32(param.MaxBatches) * .8)
	param.Steps[1] = int(float32(param.MaxBatches) * .9)
	adjust := func(value int) int {
		mod := value % 32
		return value - mod
	}
	// 调整配置文件中输入大小是 32 的整数倍
	param.Width = adjust(param.Width)
	param.Height = adjust(param.Height)

	//param.Width = int(math.Min(float64(param.Width), float64(param.Height)))
	//param.Height = param.Width

	param.Filters = (param.Classes + 5) * 3
}

const (
	YoloV4 int = iota
	YoloV4Tiny
)

type CmdLine struct {
	Path        string
	NetType     int
	Batch       int
	Subdivision int
	Classes     int
	Width       int
	Height      int
	UpdateParam int
	Absolute    int
}

type ObjectData struct {
	Classes    int
	TrainPath  string
	ValidPath  string
	NamesPath  string
	BackupPath string
}

// 全局变量 命令行参数
var cmd CmdLine

// 全局变量 使用 go:embed 嵌入资源到二进制程序中

//go:embed asset/*
var Resources embed.FS

func init() {
	// 设置命令行参数项
	flag.StringVar(&cmd.Path, "path", "./temp", "project path")
	flag.IntVar(&cmd.NetType, "type", 0, "type enum [0:yolov4,1:yolov4-tiny] default value 0")
	flag.IntVar(&cmd.Batch, "batch", 3, "batch default value 64")
	flag.IntVar(&cmd.Subdivision, "subdivision", 3, "subdivision default value 16")
	flag.IntVar(&cmd.Classes, "classes", 3, "classes default value 3")
	flag.IntVar(&cmd.Width, "width", 608, "width default value 608")
	flag.IntVar(&cmd.Height, "height", 608, "height default value 608")
	flag.IntVar(&cmd.UpdateParam, "update", 0, "only update param default 0")
	flag.IntVar(&cmd.Absolute, "absolute", 0, "use absolute path default 0")
}

// 检查错误
func checkIfError(err error) {
	if err != nil {
		panic(err)
	}
}

// 创建项目其他相关目录和文件
func CreateProject(cmd CmdLine) {
	err := os.MkdirAll(fmt.Sprint(cmd.Path, "/backup"), os.ModePerm)
	checkIfError(err)
	err = os.MkdirAll(fmt.Sprint(cmd.Path, "/data/obj"), os.ModePerm)
	checkIfError(err)

	names := fmt.Sprint(cmd.Path, "/data/obj.names")
	train := fmt.Sprint(cmd.Path, "/data/train.txt")
	valid := fmt.Sprint(cmd.Path, "/data/valid.txt")
	object := fmt.Sprint(cmd.Path, "/data/obj.data")
	readme := fmt.Sprint(cmd.Path, "/readme.txt")

	var objectParam ObjectData
	if cmd.Absolute == 1 {
		objectParam = ObjectData{
			Classes:    cmd.Classes,
			TrainPath:  train,
			ValidPath:  train,
			NamesPath:  names,
			BackupPath: fmt.Sprint(cmd.Path, "/backup"),
		}
	} else {
		objectParam = ObjectData{
			Classes:    cmd.Classes,
			TrainPath:  "data/train.txt",
			ValidPath:  "data/test.txt",
			NamesPath:  "data/obj.names",
			BackupPath: "backup/",
		}
	}

	f, err := os.Create(names)
	checkIfError(err)
	// 根据类别数量写入对应标签
	for i := 0; i < cmd.Classes; i++ {
		_, err := f.WriteString(strconv.Itoa(i))
		checkIfError(err)
		_, err = f.WriteString("\n")
		checkIfError(err)
	}

	_, err = os.Create(train)
	checkIfError(err)
	_, err = os.Create(valid)
	checkIfError(err)

	tp, err := template.ParseFS(Resources, "asset/obj.data")
	checkIfError(err)
	fs, err := os.Create(object)
	checkIfError(err)
	err = tp.Execute(fs, objectParam)
	checkIfError(err)

	in, err := Resources.Open("asset/ReadMe.md")
	checkIfError(err)
	out, err := os.OpenFile(readme, os.O_TRUNC|os.O_WRONLY|os.O_CREATE, os.ModePerm)
	checkIfError(err)
	// 复制文件
	_, err = io.Copy(out, in)
	checkIfError(err)
}

// go build -ldflags -w main.go
func main() {
	// 解析命令行参数
	flag.Parse()
	if !flag.Parsed() {
		println("please input *.exe -h get help.")
		os.Exit(0)
	}

	// 构造一个 yolov4 默认参数
	data := NewDefaultYoloV4Param()

	// 命令行参数填充到 yolov4 参数中
	data.Batch = cmd.Batch
	data.Subdivision = cmd.Subdivision
	data.Classes = cmd.Classes
	data.Width = cmd.Width
	data.Height = cmd.Height

	// 根据填充参数更新相关参数
	data.Update()

	// 创建所需要的目录
	err := os.MkdirAll(cmd.Path, os.ModePerm)
	checkIfError(err)

	if cmd.NetType == YoloV4 {
		// 使用模板和参数 data 创建 yolov4.cfg 文件
		tp, err := template.ParseFS(Resources, "asset/yolov4-custom.cfg")
		checkIfError(err)
		fs, err := os.Create(fmt.Sprint(cmd.Path, "/yolov4.cfg"))
		checkIfError(err)
		err = tp.Execute(fs, data)
		checkIfError(err)
		if cmd.UpdateParam != 1 {
			CreateProject(cmd)
		}
	} else if cmd.NetType == YoloV4Tiny {
		// 使用模板和参数 data 创建 yolo4-tiny.cfg 文件
		tp, err := template.ParseFS(Resources, "asset/yolov4-tiny-custom.cfg")
		checkIfError(err)
		fs, _ := os.Create(fmt.Sprint(cmd.Path, "/yolov4-tiny.cfg"))
		err = tp.Execute(fs, data)
		checkIfError(err)
		if cmd.UpdateParam != 1 {
			CreateProject(cmd)
		}
	}
}
